# -*- encoding: utf-8 -*-
'''
@File    :   cwru.py
@Time    :   2022/2/11 13:46
@Author  :   ZhangChaoYang
@Desc    :   美国凯斯西储大学轴承数据 https://engineering.case.edu/bearingdatacenter/download-data-file
'''

import os, sys

import numpy as np

sys.path.insert(0, os.getcwd())
from scipy.io import loadmat
from util.work_flow import trans_collection, transform

fs = 12e3  # 采样频率12kHz
rpm = 1797
class2file = {'normal': 'normal', 'ball': 'B', 'inner': 'I', 'outer': 'O'}
WC = [str(rpm)]
corpus_dir = 'E:\\zcy\\数据集\\CWRU'
corpus_name = "cwru"


def gen_corpus(trans, dim):
    if trans not in trans_collection:
        raise Exception("unsupported transformation method: {}".format(trans))

    cls_ndarray = {k: None for k in class2file}
    for _, prefix in class2file.items():
        for file in os.listdir(corpus_dir):
            if not file.startswith(prefix):
                continue
            if not file.endswith(".mat"):
                continue
            if file[-5] == '0':
                cls = 'normal'
            elif file[-5] == '1':
                cls = 'ball'
            elif file[-5] == '2':
                cls = 'inner'
            elif file[-5] == '3':
                cls = 'outer'
            ndarray = cls_ndarray[cls]
            infile = os.path.join(corpus_dir, file)
            dic = loadmat(infile)
            for key in dic.keys():  # ['__header__', '__version__', '__globals__', 'X097_DE_time', 'X097_FE_time', 'X097RPM']
                if key.count("_DE_") > 0:  # 只使用驱动端的数据
                    array = dic[key][:, 0]
                    x = transform(trans, array, 1024, 1, fs)
                    if ndarray is None:
                        ndarray = x
                    else:
                        ndarray = np.vstack((ndarray, x))
                    cls_ndarray[cls] = ndarray

    for cls, ndarray in cls_ndarray.items():
        print(cls, rpm, ndarray.shape)
        out_path = os.path.join("corpus", "multi_class", dim, corpus_name, trans)
        if not os.path.exists(out_path):
            os.makedirs(out_path)
        np.save(os.path.join(out_path, cls), ndarray)


if __name__ == '__main__':
    gen_corpus('original', '1d')
    gen_corpus('fft', '1d')
    gen_corpus('stat', '1d')
    gen_corpus('sfft', '2d')
    # gen_corpus('cwt', '2d')
